import cv2
from ultralytics import YOLO

class YOLOUtils:
    def __init__(self):
        self.model_path = 'models/qr_code_96.onnx'
        self.model = None
    
    # 载入模型 如果需要使用新模型可以先修改模型路径再直接使用即可
    def load_model(self):
        if self.model is None:
            self.model = YOLO(self.model_path)
        return self.model
    
    # 预测
    def predict(self, img):
        frame = cv2.resize(img, (640, 640))
        return self.load_model(frame)
